{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Relax 构建\n",
    "\n",
    "本教程将演示如何创建 Relax 函数和程序。\n",
    "\n",
    "介绍定义 Relax 函数的多种方法，包括使用 TVMScript 和 Relax NNModule API。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 TVMScript 创建 Relax 程序\n",
    "\n",
    "TVMScript 是一种领域特定语言，用于表示 Apache TVM 的中间表示（IR）。它是一种 Python 方言，可用于定义包含 TensorIR 和 Relax 函数的 IRModule。\n",
    "\n",
    "在本节中，将展示如何使用 TVMScript 仅通过高级 Relax 运算符来定义简单的多层感知机（MLP）模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(data: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), w0: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), b0: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), w1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), b1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(w0, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            lv1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(data, lv, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            lv0: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv1, b0)\n",
       "            lv1_1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(lv0)\n",
       "            lv4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(w1, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            lv5: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(lv1_1, lv4, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            lv2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv5, b1)\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(lv2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> lv2\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tvm import relax, topi\n",
    "from tvm.script import ir as I\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "\n",
    "\n",
    "@I.ir_module\n",
    "class RelaxModule:\n",
    "    @R.function\n",
    "    def forward(\n",
    "        data: R.Tensor((\"n\", 784), dtype=\"float32\"),\n",
    "        w0: R.Tensor((128, 784), dtype=\"float32\"),\n",
    "        b0: R.Tensor((128,), dtype=\"float32\"),\n",
    "        w1: R.Tensor((10, 128), dtype=\"float32\"),\n",
    "        b1: R.Tensor((10,), dtype=\"float32\"),\n",
    "    ) -> R.Tensor((\"n\", 10), dtype=\"float32\"):\n",
    "        with R.dataflow():\n",
    "            lv0 = R.matmul(data, R.permute_dims(w0)) + b0\n",
    "            lv1 = R.nn.relu(lv0)\n",
    "            lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1\n",
    "            R.output(lv2)\n",
    "        return lv2\n",
    "\n",
    "\n",
    "RelaxModule.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Relax 不仅是一种图级别的中间表示（IR），还支持跨级别的表示与转换。具体来说，可以在 Relax 函数中直接调用 TensorIR 函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">relu</span>(x: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, y: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        n, m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        X <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(x, (n, m))\n",
       "        Y <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(y, (n, m))\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, m):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;relu&quot;</span>):\n",
       "                vi, vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i, j])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(X[vi, vj])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
       "                Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(X[vi, vj], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
       "\n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(data: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), w0: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), b0: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), w1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), b1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        cls <span style=\"color: #AA22FF; font-weight: bold\">=</span> Module\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(w0, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            lv1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(data, lv, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            lv0: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv1, b0)\n",
       "            lv1_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu, (lv0,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(w1, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            lv5: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(lv1_1, lv4, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            lv2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv5, b1)\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(lv2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> lv2\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@I.ir_module\n",
    "class RelaxModuleWithTIR:\n",
    "    @T.prim_func\n",
    "    def relu(x: T.handle, y: T.handle):\n",
    "        n, m = T.int64(), T.int64()\n",
    "        X = T.match_buffer(x, (n, m), \"float32\")\n",
    "        Y = T.match_buffer(y, (n, m), \"float32\")\n",
    "        for i, j in T.grid(n, m):\n",
    "            with T.block(\"relu\"):\n",
    "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))\n",
    "\n",
    "    @R.function\n",
    "    def forward(\n",
    "        data: R.Tensor((\"n\", 784), dtype=\"float32\"),\n",
    "        w0: R.Tensor((128, 784), dtype=\"float32\"),\n",
    "        b0: R.Tensor((128,), dtype=\"float32\"),\n",
    "        w1: R.Tensor((10, 128), dtype=\"float32\"),\n",
    "        b1: R.Tensor((10,), dtype=\"float32\"),\n",
    "    ) -> R.Tensor((\"n\", 10), dtype=\"float32\"):\n",
    "        n = T.int64()\n",
    "        cls = RelaxModuleWithTIR\n",
    "        with R.dataflow():\n",
    "            lv0 = R.matmul(data, R.permute_dims(w0)) + b0\n",
    "            lv1 = R.call_tir(cls.relu, lv0, R.Tensor((n, 128), dtype=\"float32\"))\n",
    "            lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1\n",
    "            R.output(lv2)\n",
    "        return lv2\n",
    "\n",
    "\n",
    "RelaxModuleWithTIR.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{note}\n",
    ":class: alert alert-info\n",
    "\n",
    "你可能会注意到打印输出的内容与编写的 TVMScript 代码有所不同。这是因为以标准格式打印 IRModule，同时支持输入时的语法糖。\n",
    "\n",
    "例如，可以将多个操作合并到一行中，如下所示：\n",
    "\n",
    "```python\n",
    "lv0 = R.matmul(data, R.permute_dims(w0)) + b0\n",
    "```\n",
    "\n",
    "然而，规范化表达式要求每个绑定中只能有一个操作。因此，打印输出的内容与编写的 TVMScript 代码不同，如下所示：\n",
    "\n",
    "```python\n",
    "lv: R.Tensor((784, 128), dtype=\"float32\") = R.permute_dims(w0, axes=None)\n",
    "lv1: R.Tensor((n, 128), dtype=\"float32\") = R.matmul(data, lv, out_dtype=\"void\")\n",
    "lv0: R.Tensor((n, 128), dtype=\"float32\") = R.add(lv1, b0)</p></div>\n",
    "```\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 NNModule API 创建 Relax 程序\n",
    "\n",
    "除了 TVMScript 外，还提供了一种类似 PyTorch 的 API 来定义神经网络。它的设计更加直观且易于使用。\n",
    "\n",
    "在本节中，将展示如何使用 Relax NNModule API 来定义相同的多层感知机（MLP）模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from tvm.relax.frontend import nn\n",
    "\n",
    "\n",
    "class NNModule(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(784, 128)\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = self.relu1(x)\n",
    "        x = self.fc2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在定义完 NNModule 后，可以通过 `export_tvm` 将其导出为 TVM IRModule。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            permute_dims: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(fc1_weight, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            matmul: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(x, permute_dims, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            add: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(matmul, fc1_bias)\n",
       "            relu: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(add)\n",
       "            permute_dims1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(fc2_weight, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            matmul1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(relu, permute_dims1, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            add1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(matmul1, fc2_bias)\n",
       "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> add1\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod, params = NNModule().export_tvm({\"forward\": {\"x\": nn.spec.Tensor((\"n\", 784), \"float32\")}})\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "还可以在 NNModule 中插入自定义的函数调用，例如 Tensor Expression (TE)、TensorIR 函数或其他 TVM 打包函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">relu</span>(var_env_linear: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, var_compute: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        env_linear <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_env_linear, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
       "        compute <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_compute, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;compute&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(env_linear[v_i0, v_i1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(compute[v_i0, v_i1])\n",
       "                compute[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(env_linear[v_i0, v_i1], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
       "\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">tir_linear</span>(x: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, w: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, b: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, z: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        M, K <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        X <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(x, (M, K))\n",
       "        N <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        W <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(w, (N, K))\n",
       "        B <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(b, (N,))\n",
       "        Z <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(z, (M, N))\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(M, N, K):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;linear&quot;</span>):\n",
       "                vi, vj, vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i, j, k])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(X[vi, vk], W[vj, vk])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Z[vi, vj])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
       "                Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> X[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> W[vj, vk]\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(M, N):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;add&quot;</span>):\n",
       "                vi, vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i, j])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Z[vi, vj], B[vj])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Z[vi, vj])\n",
       "                Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[vj]\n",
       "\n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        cls <span style=\"color: #AA22FF; font-weight: bold\">=</span> Module\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            env_linear <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_dps_packed(<span style=\"color: #BA2121\">&quot;env.linear&quot;</span>, (x, fc1_weight, fc1_bias), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu, (env_linear,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>tir_linear, (lv, fc2_weight, fc2_bias), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv1\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@T.prim_func\n",
    "def tir_linear(x: T.handle, w: T.handle, b: T.handle, z: T.handle):\n",
    "    M, N, K = T.int64(), T.int64(), T.int64()\n",
    "    X = T.match_buffer(x, (M, K), \"float32\")\n",
    "    W = T.match_buffer(w, (N, K), \"float32\")\n",
    "    B = T.match_buffer(b, (N,), \"float32\")\n",
    "    Z = T.match_buffer(z, (M, N), \"float32\")\n",
    "    for i, j, k in T.grid(M, N, K):\n",
    "        with T.block(\"linear\"):\n",
    "            vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "            with T.init():\n",
    "                Z[vi, vj] = 0\n",
    "            Z[vi, vj] = Z[vi, vj] + X[vi, vk] * W[vj, vk]\n",
    "    for i, j in T.grid(M, N):\n",
    "        with T.block(\"add\"):\n",
    "            vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "            Z[vi, vj] = Z[vi, vj] + B[vj]\n",
    "\n",
    "\n",
    "class NNModuleWithTIR(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(784, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        n = x.shape[0]\n",
    "        # We can call external functions using nn.extern\n",
    "        x = nn.extern(\n",
    "            \"env.linear\",\n",
    "            [x, self.fc1.weight, self.fc1.bias],\n",
    "            out=nn.Tensor.placeholder((n, 128), \"float32\"),\n",
    "        )\n",
    "        # We can also call TensorIR via Tensor Expression API in TOPI\n",
    "        x = nn.tensor_expr_op(topi.nn.relu, \"relu\", [x])\n",
    "        # We can also call other TVM packed functions\n",
    "        x = nn.tensor_ir_op(\n",
    "            tir_linear,\n",
    "            \"tir_linear\",\n",
    "            [x, self.fc2.weight, self.fc2.bias],\n",
    "            out=nn.Tensor.placeholder((n, 10), \"float32\"),\n",
    "        )\n",
    "        return x\n",
    "\n",
    "\n",
    "mod, params = NNModuleWithTIR().export_tvm(\n",
    "    {\"forward\": {\"x\": nn.spec.Tensor((\"n\", 784), \"float32\")}}\n",
    ")\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 Block Builder API 创建 Relax 程序\n",
    "\n",
    "除了上述 API 外，我们还提供了 Block Builder API 用于创建 Relax 程序。它是一种 IR 构建器 API，更加底层，广泛用于 TVM 的内部逻辑中，例如编写自定义的 Pass。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;v&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;v&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        v <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(fc1_weight, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            lv1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(x, lv, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            lv2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv1, fc1_bias)\n",
       "            lv3: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(lv2)\n",
       "            lv4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(fc2_weight, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            lv5: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(lv3, lv4, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            lv6: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv5, fc2_bias)\n",
       "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv6\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> lv6\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bb = relax.BlockBuilder()\n",
    "n = T.int64()\n",
    "x = relax.Var(\"x\", R.Tensor((n, 784), \"float32\"))\n",
    "fc1_weight = relax.Var(\"fc1_weight\", R.Tensor((128, 784), \"float32\"))\n",
    "fc1_bias = relax.Var(\"fc1_bias\", R.Tensor((128,), \"float32\"))\n",
    "fc2_weight = relax.Var(\"fc2_weight\", R.Tensor((10, 128), \"float32\"))\n",
    "fc2_bias = relax.Var(\"fc2_bias\", R.Tensor((10,), \"float32\"))\n",
    "with bb.function(\"forward\", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):\n",
    "    with bb.dataflow():\n",
    "        lv0 = bb.emit(relax.op.matmul(x, relax.op.permute_dims(fc1_weight)) + fc1_bias)\n",
    "        lv1 = bb.emit(relax.op.nn.relu(lv0))\n",
    "        gv = bb.emit(relax.op.matmul(lv1, relax.op.permute_dims(fc2_weight)) + fc2_bias)\n",
    "        bb.emit_output(gv)\n",
    "    bb.emit_func_output(gv)\n",
    "\n",
    "mod = bb.get()\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此外，Block Builder API 支持构建包含 Relax 函数、TensorIR 函数以及其他 TVM 打包函数的跨级别 IRModule。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "tags": [
     "hide-output",
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">relu</span>(var_lv: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, var_compute: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        v <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        lv <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_lv, (v, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
       "        compute <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_compute, (v, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(v, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;compute&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(lv[v_i0, v_i1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(compute[v_i0, v_i1])\n",
       "                compute[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(lv[v_i0, v_i1], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
       "\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">tir_linear</span>(x: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, w: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, b: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, z: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        M, K <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        X <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(x, (M, K))\n",
       "        N <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        W <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(w, (N, K))\n",
       "        B <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(b, (N,))\n",
       "        Z <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(z, (M, N))\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(M, N, K):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;linear&quot;</span>):\n",
       "                vi, vj, vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i, j, k])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(X[vi, vk], W[vj, vk])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Z[vi, vj])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
       "                Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> X[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> W[vj, vk]\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(M, N):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;add&quot;</span>):\n",
       "                vi, vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i, j])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Z[vi, vj], B[vj])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Z[vi, vj])\n",
       "                Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Z[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[vj]\n",
       "\n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;v&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;v&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        v <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
       "        cls <span style=\"color: #AA22FF; font-weight: bold\">=</span> Module\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_dps_packed(<span style=\"color: #BA2121\">&quot;env.linear&quot;</span>, (x, fc1_weight, fc1_bias), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu, (lv,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>tir_linear, (lv1, fc2_weight, fc2_bias), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((v, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv2\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> lv2\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bb = relax.BlockBuilder()\n",
    "with bb.function(\"forward\", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):\n",
    "    with bb.dataflow():\n",
    "        lv0 = bb.emit(\n",
    "            relax.call_dps_packed(\n",
    "                \"env.linear\",\n",
    "                [x, fc1_weight, fc1_bias],\n",
    "                out_sinfo=relax.TensorStructInfo((n, 128), \"float32\"),\n",
    "            )\n",
    "        )\n",
    "        lv1 = bb.emit_te(topi.nn.relu, lv0)\n",
    "        tir_gv = bb.add_func(tir_linear, \"tir_linear\")\n",
    "        gv = bb.emit(\n",
    "            relax.call_tir(\n",
    "                tir_gv,\n",
    "                [lv1, fc2_weight, fc2_bias],\n",
    "                out_sinfo=relax.TensorStructInfo((n, 10), \"float32\"),\n",
    "            )\n",
    "        )\n",
    "        bb.emit_output(gv)\n",
    "    bb.emit_func_output(gv)\n",
    "mod = bb.get()\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "需要注意的是，Block Builder API 不如上述 API 那样用户友好，但它是与 IR 定义紧密结合的最低层级 API。对于仅希望定义和转换机器学习模型的用户，推荐使用上述 API。然而，对于那些希望构建更复杂转换的用户，Block Builder API 则是更为灵活的选择。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结\n",
    "本教程展示了如何针对不同的使用场景，通过 TVMScript、NNModule API、Block Builder API 以及 PackedFunc API 来创建 Relax 程序。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
